#include <stdio.h>
#include <math.h>
#include <stdlib.h>
#include <string.h>

#include <icommon.h>
#ifdef DMALLOC
#include <imp/dmalloc.h>
#endif
struct watermark {
	int    pwm1;
	int    pwm2;
	int    wcnt;
	float strength;
	float relation;
    float *srcbuf;
    float *dstbuf;
    float *midbuf;
    float *resbuf;
};

static float watermarkn[64] = { /* ingenic */
	126.000000,	 24.000000,	 24.000000,	 24.000000,	 24.000000,	 24.000000,	 24.000000,	126.000000,
	 62.000000,	127.000000,	102.000000,	102.000000,	102.000000,	102.000000,	102.000000,	248.000000,
	254.000000,	230.000000,	118.000000,	 62.000000,	126.000000,	 67.000000,	127.000000,	 60.000000,
	 24.000000,	126.000000,	102.000000,	102.000000,	127.000000,	  6.000000,	102.000000,	126.000000,
	 62.000000,	127.000000,	102.000000,	102.000000,	102.000000,	102.000000,	102.000000,	255.000000,
	 24.000000,	  0.000000,	 28.000000,	 30.000000,	 24.000000,	 24.000000,	 24.000000,	126.000000,
	 24.000000,	126.000000,	102.000000,	103.000000,	  3.000000,	 71.000000,	102.000000,	126.000000,
};

static float coscoff[8][8] = {
    {0.353553,  0.353553,  0.353553,  0.353553,  0.353553,  0.353553,  0.353553,  0.353553},
    {0.490393,  0.415735,  0.277785,  0.097545, -0.097545, -0.277785, -0.415735, -0.490393},
    {0.461940,  0.191341, -0.191341, -0.461940, -0.461940, -0.191342,  0.191341,  0.461940},
    {0.415735, -0.097545, -0.490393, -0.277785,  0.277785,  0.490393,  0.097545, -0.415735},
    {0.353554, -0.353554, -0.353554,  0.353554,  0.353554, -0.353554, -0.353554,  0.353554},
    {0.277785, -0.490393,  0.097545,  0.415735, -0.415735, -0.097545,  0.490393, -0.277785},
    {0.191341, -0.461940,  0.461940, -0.191341, -0.191342,  0.461940, -0.461940,  0.191341},
    {0.097545, -0.277785,  0.415735, -0.490393,  0.490393, -0.415735,  0.277785, -0.097545}
};

static float mindct[64] = {
	74.999832,	  7.334908,	 -6.836268,	 -6.089550,	  6.000004,	  5.526941,	 -5.031826,	 -4.745100,
	-4.475156,	  4.379658,	 -4.129858,	 -4.037567,	 -3.963046,	 -3.781809,	 -3.637949,	  3.447896,
	3.290550,	  3.261708,	 -3.098304,	  3.080270,	 -2.846243,	  2.840293,	 -2.772887,	 -2.767764,
	-2.636348,	  2.483991,	 -2.473889,	  2.414223,	 -2.307114,	  2.301600,	 -2.044400,	 -1.902556,
	-1.736412,	 -1.689251,	 -1.566984,	  1.500001,	 -1.490217,	  1.467793,	 -1.465076,	 -1.457760,
	-1.344692,	  1.012724,	  1.000003,	 -0.952022,	  0.874869,	  0.857513,	  0.821410,	  0.767766,
	0.672561,	 -0.669244,	 -0.668013,	 -0.599824,	 -0.544652,	  0.426042,	  0.414220,	  0.405658,
	0.402439,	 -0.318998,	 -0.285764,	  0.234587,	 -0.219605,	 -0.212240,	 -0.057854,	  0.000293,
};

static void fabs64(float *dst, float *src)
{
	int i = 0;
	for (i = 0; i < 64; i++) {
		dst[i] = fabs(src[i]);
	}
}

static void fsort64(float *s, int l, int r)
{
    if (l < r)
    {
        //Swap(s[l], s[(l + r) / 2]); //将中间的这个数和第一个数交换 参见注1
        int i = l, j = r;
		float x = s[l];
        while (i < j)
        {
            while((i < j) && (s[j] <= x)) // 从右向左找第一个小于x的数
                j--;
            if(i < j)
                s[i++] = s[j];

            while((i < j) && (s[i] > x)) // 从左向右找第一个大于等于x的数
                i++;
            if(i < j)
                s[j--] = s[i];
        }
        s[i] = x;
        fsort64(s, l, i - 1); // 递归调用
        fsort64(s, i + 1, r);
    }
}

static int fpos64(float *data, float a)
{
	int i = 0;

	for (i = 0; i < 64; i++) {
		if (fabs(a - data[i]) < 0.000005) {
			return i;
		}
	}
	return -1;
}

float corcoeff8(float *x, float *y, int n)
{
	int i = 0;
	float xa = 0.0;
	float cosxy = 0.0, xx = 0.0, yy = 0.0;

	for (i = 0; i < n; i++) {
		xa += x[i];
	}
	xa = xa / n;

	for (i = 0; i < n; i++) {
		cosxy += (x[i] - xa) * (y[i] - xa);
		xx += (x[i] - xa) * (x[i] - xa);
		yy += (y[i] - xa) * (y[i] - xa);
	}
	return fabs(cosxy / (sqrt(xx) * sqrt(yy)));
}

#if 0
static void watermark_debug_matrix8(float *matrix, int n)
{
    int i = 0;

    for (i = 0; i < n; i++) {
        printf("%f,\t", matrix[i]);
        if ((i + 1) % 8 == 0) {
            printf("\n");
        }
    }
}
#endif

static void watermark_dct8(float *F, float *f, float *dTemp, int N)
{
    int n, m, x;

    memset(F, 0, sizeof(float) * N * N);
    memset(dTemp, 0, sizeof(float) * N * N);

    //一维变换
    for (m = 0; m < N; m++) {
        for (x = 0; x < N; x++) {
            for (n = 0; n < N; n++) {
                dTemp[m * N + n] += f[x * N + n] * coscoff[m][x];
            }
        }
    }
    //第二次一维变换
    for (n = 0; n < N; n++) {
        for (x = 0; x < N; x++) {
            for (m = 0; m < N; m++) {
                F[m * N + n] += dTemp[m * N + x] * coscoff[n][x];
            }
        }
    }
}

static void watermark_idct8(float *f, float *F, float *dTemp, int N)
{
    int m,y,x;

    memset(f, 0, sizeof(float) * N * N);
    memset(dTemp, 0, sizeof(float) * N * N);

    //一维变换
    for (m = 0; m < N; m++) {
        for (y = 0; y < N; y++) {
            for (x = 0; x < N; x++) {
                dTemp[x * N + y] += F[x * N + m] * coscoff[m][y];
            }
        }
    }
    //第二次一维变换
    for (m = 0; m < N; m++) {
        for (x = 0; x < N; x++) {
            for (y = 0; y < N; y++) {
                f[x * N + y] += dTemp[m * N + y] * coscoff[m][x];
            }
        }
    }
}

watermark_t *watermark_init(void)
{
    watermark_t *handler = NULL;

    handler = malloc(sizeof(watermark_t) + sizeof(float) * 256);
    if (handler == NULL) {
        c_log(C_LOG_ERROR, "calloc wm failed\n");
        goto err_calloc_watermark;
    }

	handler->pwm1 = 33;
	handler->pwm2 = 12;
	handler->wcnt = 16;
	handler->strength = 0.1;
	handler->relation = 0.42;
    handler->srcbuf = (float *)((char *)handler + sizeof(watermark_t));
    handler->dstbuf = (float *)((char *)(handler->srcbuf) + sizeof(float) * 64);
    handler->midbuf = (float *)((char *)(handler->dstbuf) + sizeof(float) * 64);
    handler->resbuf = (float *)((char *)(handler->midbuf) + sizeof(float) * 64);

    return handler;

err_calloc_watermark:
    return NULL;
}

void watermark_deinit(watermark_t *handler)
{
    if (handler) {
        free(handler);
    }
}

static void watermark_fetch_matrix8_bw8(float *matrix, unsigned char *buf, int w, int h)
{
    int i = 0;

    for (i = 0; i < 64; i++) {
        matrix[i] = (unsigned int)(buf[i / 8 * w + i % 8]);
    }
}

static void watermark_fill_matrix8_bw8(unsigned char *buf, float *matrix, int w, int h)
{
    int i = 0;
    for (i = 0; i < 64; i++) {
        buf[i / 8 * w + i % 8] = (unsigned int)(matrix[i] + 0.5);
    }
}

int embed_watermark(watermark_t *handler, unsigned char *ybuf, int w, int h)
{
    int p = 0, q = 0;
    unsigned char *buf = NULL;
	int p30 = 0, p31 = 0, p32 = 0, spwm1 = 1, spwm2 = 1, s30 = 1, s31 = 1, s32 = 1;

    if (w < 64 || h < handler->wcnt) {
        c_log(C_LOG_ERROR, "wxh=(%d,%d) is too small\n", w, h);
        return -1;
    }
    for (p = 0; p < handler->wcnt; p += 8) {
        for (q = 0; q < 64; q += 8) {
			buf = ybuf + p * w + q;
			watermark_fetch_matrix8_bw8(handler->srcbuf, buf, w, h);

			watermark_dct8(handler->dstbuf, handler->srcbuf, handler->midbuf, 8);

			if (fabs(handler->dstbuf[0]) < 80) {
				memcpy(handler->dstbuf, mindct, sizeof(float) * 64);
			}

			fabs64(handler->midbuf, handler->dstbuf);
			memcpy(handler->resbuf, handler->midbuf, sizeof(float) * 64);

			fsort64(handler->midbuf, 0, 63);

			p30 = fpos64(handler->resbuf, handler->midbuf[30]);
			p31 = fpos64(handler->resbuf, handler->midbuf[31]);
			p32 = fpos64(handler->resbuf, handler->midbuf[32]);
			spwm1 = handler->dstbuf[handler->pwm1] > 0 ? 1 : -1;
			spwm2 = handler->dstbuf[handler->pwm2] > 0 ? 1 : -1;
			s30 = handler->dstbuf[p30] > 0 ? 1 : -1;
			s31 = handler->dstbuf[p31] > 0 ? 1 : -1;
			s32 = handler->dstbuf[p32] > 0 ? 1 : -1;
			handler->dstbuf[p31] = fabs(handler->dstbuf[p31]) > 0.5 ? handler->dstbuf[p31] : s31 * 0.5;
			handler->dstbuf[p30] = s30 * fabs(handler->dstbuf[p31]);
			handler->dstbuf[p32] = s32 * fabs(handler->dstbuf[p31]);

			handler->dstbuf[handler->pwm1] = spwm1 * fabs(handler->dstbuf[p32]) * (1 + handler->strength * watermarkn[p / 8 * 16 + (q / 8) * 2 + 0]);
			handler->dstbuf[handler->pwm2] = spwm2 * fabs(handler->dstbuf[p32]) * (1 + handler->strength * watermarkn[p / 8 * 16 + (q / 8) * 2 + 1]);


			watermark_idct8(handler->srcbuf, handler->dstbuf, handler->midbuf, 8);

			watermark_fill_matrix8_bw8(buf, handler->srcbuf, w, h);
        }
    }
    return 0;
}

int check_watermark(watermark_t *handler, unsigned char *ybuf, int w, int h)
{
    int p = 0, q = 0;
    unsigned char *buf = NULL;
	float wmbuf[64], rrate;

    if (w < 64 || h < handler->wcnt) {
        c_log(C_LOG_ERROR, "wxh=(%d,%d) is too small\n", w, h);
        return -1;
    }

	memset(wmbuf, 0, sizeof(float));
	for (p = 0; p < handler->wcnt; p += 8) {
		for (q = 0; q < 64; q += 8) {
			buf = ybuf + p * w + q;
			watermark_fetch_matrix8_bw8(handler->srcbuf, buf, w, h);

			watermark_dct8(handler->dstbuf, handler->srcbuf, handler->midbuf, 8);

			fabs64(handler->midbuf, handler->dstbuf);
			fsort64(handler->midbuf, 0, 63);
			handler->midbuf[32] = handler->midbuf[32] != 0.0 ? handler->midbuf[32] : 0.00005;
			wmbuf[(p / 8) * 16 + (q / 8) * 2 + 0] = ((fabs(handler->dstbuf[handler->pwm1]) / handler->midbuf[32]) - 1) / handler->strength;
			wmbuf[(p / 8) * 16 + (q / 8) * 2 + 1] = ((fabs(handler->dstbuf[handler->pwm2]) / handler->midbuf[32]) - 1) / handler->strength;
		}
	}

	rrate = corcoeff8(watermarkn, wmbuf, handler->wcnt * 2);

	return rrate >= handler->relation ? 1 : 0;
}
